"""
@Filename       : metric.py
@Create Time    : 2020/11/2 20:00
@Author         : Rylynn
@Description    : 

"""
import numpy as np
import torch
from tqdm import tqdm


def hits_k(probs, true_nodes, k=10):
    count = 0
    hit = 0
    for prob, true_node in zip(probs, true_nodes):
        count = count + 1
        top_k = prob.argsort()[-k:][::-1]
        if true_node in top_k:
            hit = hit + 1

    return hit, count


def apk(true_nodes, predicted, k=10):
    if len(predicted) > k:
        predicted = predicted[:k]

    score = 0.0
    num_hits = 0.0
    for i, p in enumerate(predicted):
        if p in true_nodes and p not in predicted[:i]:
            num_hits += 1.0
            score += num_hits / (i + 1.0)

    if not true_nodes:
        return 0.0
    return score / min(len(true_nodes), k)


def map_k(probs, true_nodes, k=10):
    count = len(true_nodes)

    predicted = [np.argsort(p_)[-k:][::-1] for p_ in probs]
    actual = [[y_] for y_ in true_nodes]

    ap = np.sum(apk(a, p, k) for a, p in zip(actual, predicted))
    return ap, count


def run_evaluation(model, dataloader):
    model.eval()
    result = {'hits@10': [],
              'hits@50': [],
              'hits@100': [],
              'map@10': [],
              'map@50': [],
              'map@100': []}

    for (data, data_length) in tqdm(dataloader):
        # print(data)
        data = data.cuda()
        prob = model(data, data_length)
        prob = prob.cpu().detach().numpy()
        next_node = data[:, 1:].contiguous().view(-1)
        next_node = next_node.tolist()

        prob_y = []
        y = []
        for p, n in zip(prob, next_node):
            if n != 0:
                prob_y.append(p)
                y.append(n)
        prob_y = np.array(prob_y)

        result['hits@10'].append(hits_k(prob_y, y, 10))
        result['hits@50'].append(hits_k(prob_y, y, 50))
        result['hits@100'].append(hits_k(prob_y, y, 100))
        result['map@10'].append(map_k(prob_y, y, 10))
        result['map@50'].append(map_k(prob_y, y, 50))
        result['map@100'].append(map_k(prob_y, y, 100))

    for k, v in result.items():
        v0_sum = sum([item[0] for item in v])
        v1_sum = sum([item[1] for item in v])
        result[k] = v0_sum / v1_sum

    print(result)
    return result
